# !/usr/bin/env python
# -*- coding: utf-8 -*-
from pyspark.sql import SparkSession
from pyspark.sql.functions import isnull, when, count, col
from pyspark.ml.feature import VectorAssembler, StringIndexer
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.regression import LinearRegression, RandomForestRegressor
from pyspark.ml.evaluation import MulticlassClassificationEvaluator

# https://www.kaggle.com/uciml/pima-indians-diabetes-database?ref=hackernoon.com
from pyspark.sql.types import DoubleType, StringType

'''
任务类型: 随机森林回归主要用于回归任务。在回归任务中, 算法试图预测一个连续的数值输出, 而不是一个离散的类别。
输出: 随机森林回归的输出是一个连续的数值, 表示输入数据的预测结果。
算法原理: 随机森林回归同样基于决策树, 但在回归任务中, 每个决策树的输出是一个实数值。最终的预测结果是多个决策树输出的平均值或加权平均值。

字段说明: SepalLength(花萼长度), SepalWidth(花萼宽度), PetalLength(花瓣长度), PetalWidth(花瓣宽度), Species(品种).

品种说明: Setosa(山鸢尾), Versicolor(变色鸢尾), Virginical(维吉尼亚鸢尾).
'''
'''
准确率(accuracy)、精确率(precision)、召回率(recall)、F1-Score、ROC曲线、P-R曲线、AUC面积
https://blog.csdn.net/michael_f2008/article/details/104938824
'''


def get_logistic(name, labelCol, featuresCol):
    if name == 'Logistic':
        # 逻辑回归
        return LogisticRegression(labelCol=labelCol, featuresCol=featuresCol, maxIter=10)
    if name == 'RandomForest':
        # 随机森林回归
        return RandomForestRegressor(labelCol=labelCol, featuresCol=featuresCol, maxDepth=10)
    if name == 'Linear':
        # 线性回归回归
        return LinearRegression(labelCol=labelCol, featuresCol=featuresCol, maxIter=10)
    return None


if __name__ == '__main__':
    spark = SparkSession.builder.appName("logistic_learn").master("local[*]").getOrCreate()

    # 加载数据
    df = spark.read.option("header", "true").option("encoding", "utf-8").csv("../data/iris.csv")
    df.printSchema()

    # 缺失数据处理
    df.select([count(when(isnull(c), c)).alias(c) for c in df.columns]).show()

    # 转换 Spark 中 DateFrame 数据类型。
    df = df.withColumn("SepalLength", col("SepalLength").cast(DoubleType()))
    df = df.withColumn("SepalWidth", col("SepalWidth").cast(DoubleType()))
    df = df.withColumn("PetalLength", col("PetalLength").cast(DoubleType()))
    df = df.withColumn("PetalWidth", col("PetalWidth").cast(DoubleType()))
    df = df.withColumn("Species", col("Species").cast(StringType()))

    # 输出 Spark 中 DataFrame 字段和数据类型。
    print("Changed Spark DataFrame Data Type:")
    df.printSchema()

    # 使用 StringIndexer 转换 Species 列。
    indexer = StringIndexer(inputCol="Species", outputCol="SpeciesIdx")
    # 拟合并转换数据。
    indexDF = indexer.fit(df).transform(df)

    # 输出 StringIndexer 的转换效果。
    print("The Effect of StringIndexer:")
    indexDF.show()

    # 将 SpeciesIdx 列复制为 Label 列。
    dataset = indexDF.withColumn("Label", col("SpeciesIdx"))

    # VectorAssembler 将多个特征合并为一个特征向量。
    required_features: list = ["SepalLength", "SepalWidth", "PetalLength", "PetalWidth"]
    assembler = VectorAssembler(inputCols=required_features, outputCol="Features")
    transformed_data = assembler.transform(dataset)
    transformed_data.show()

    # 将数据随机分成训练集和测试集，并设置可重复性的种子。
    (training_data, test_data) = transformed_data.randomSplit([0.8, 0.2], seed=2020)
    print("训练数据集总数: " + str(training_data.count()))
    print("测试数据集总数: " + str(test_data.count()))

    class_arr = ["Logistic", "RandomForest", "Linear"]
    for name in class_arr:
        lr = get_logistic(name, 'Label', 'Features')
        model = lr.fit(training_data)
        # 在测试集上进行预测
        lr_predictions = model.transform(test_data)
        print(f"==================={name} prediction=====================")
        lr_predictions.show()

        #  评估随机森林分类器模型
        multi_evaluator = MulticlassClassificationEvaluator(labelCol='Label', metricName='accuracy')
        print(name + ' Logistic Regression Accuracy:', multi_evaluator.evaluate(lr_predictions))

    spark.stop()
